"""Combine testing results of the three models to get final accuracy."""
from __future__ import division
from __future__ import print_function
import argparse
import numpy as np
import ast    # Abstract Syntax Trees
import pdb


def main():
    def arg_as_list(s):
        v = ast.literal_eval(s)
        if type(v) is not list:
            raise argparse.ArgumentTypeError("Argument \"%s\" is not a list" % (s))
        for ls in v:
            if len(ls) != 3:
                raise argparse.ArgumentTypeError("len(%s) is not 3" % (str(ls)))
        return v
    parser = argparse.ArgumentParser(description="combine predictions")
    parser.add_argument('--iframe', type=str, required=True,
                        help='iframe score file.')
    parser.add_argument('--mv', type=str, required=True,
                        help='motion vector score file.')
    parser.add_argument('--res', type=str, required=True,
                        help='residual score file.')
    # [(wi, wm, wr), ]
    parser.add_argument('--combine', type=arg_as_list,
                        default=[(1, 0, 0), (0, 1, 0), (0, 0, 1)],
                        help='iframe weight.')

    args = parser.parse_args()
    print(args.combine)

    with np.load(args.iframe) as iframe:
        with np.load(args.mv) as mv:
            with np.load(args.res) as res:
                i_score = np.array([score[0][0] for score in iframe['scores']])
                mv_score = np.array([score[0][0] for score in mv['scores']]).reshape(-1, 11)
                res_score = np.array([score[0][0] for score in res['scores']])

                i_label = np.array([score[1] for score in iframe['scores']])
                mv_label = np.array([score[1] for score in mv['scores']])
                res_label = np.array([score[1] for score in res['scores']])
                assert np.alltrue(i_label == mv_label) and np.alltrue(i_label == res_label)

                ACCdown = len(mv['names'])
                ODRdown = np.count_nonzero(i_label != 0)

                for (wi, wm, wr) in args.combine:
                    print('combine', (wi, wm, wr))
                    combined_score = i_score * wi + mv_score * wm + res_score * wr
                    com_pre = np.argmax(combined_score, axis=1)

                    data = np.column_stack((mv['names'], i_label, com_pre))
                    fn = '/'.join(args.iframe.split('/')[:-1]) + \
                        '/predict_label_combine' + \
                        str(wi) + str(wm) + str(wr) + '.txt'
                    with open(fn, 'w') as f:
                        for i in range(data.shape[0]):
                            f.write(' '.join(data[i]) + '\n')
                    # np.savetxt(fn, data)  # TODO

                    ACCup, ODRup = 0, 0
                    PEA = .0
                    all_event_ave = .0
                    for clas in range(i_score.shape[1]):
                        up = np.count_nonzero((com_pre == i_label) &
                                              (i_label == clas))
                        down = np.count_nonzero(i_label == clas)
                        ACCup += up
                        all_event_ave += up/(down*11)
                        if clas != 0:
                            ODRup += up
                            PEA += up/(down*10)
                        print('{:>15}'.format('Ev' + str(clas+100) + ' acc:'), up/down)
                    print('{:>15}'.format('Pos Event Ave:'), PEA)
                    print('{:>15}'.format('ODR:'), ODRup/ODRdown)
                    print('{:>15}'.format('All Event Ave:'), all_event_ave)
                    print('{:>15}'.format('All Event Acc:'), ACCup/ACCdown)

                # accuracy = float(sum(com_pre == i_label)) / ACCdown
                # print('Accuracy: %f (%d).' % (accuracy, ACCdown))


if __name__ == '__main__':
    main()
